"""
author : DengXiuqi
date : 2018.10
email : dengxiuqi@163.com
"""

from config import *
import numpy as np
from skimage import io
import scipy.misc
import os, glob, random

MEAN_VALUES = np.array([123.68, 116.779, 103.939]).reshape((1, 1, 3))

class ImgData(object):
    def __init__(self, batchSize, dataType):
        # self.dataType = dataType
        assert dataType.lower() in ['train', 'test']
        self.cursor = 0
        self.epoch = 0
        self.batch_size = batchSize
        self.size = test_size if dataType == "test" else train_size
        self.clear_path = './data/%s_clear/' % dataType
        self.blur_path = './data/%s_blur/' % dataType
        self.file_num = list(range(self.size))
        if dataType == "train":
            random.shuffle(self.file_num)

    def next_batch(self):
        def load_batch(fileList, imgSize):
            img_batch = []
            for filename in fileList:
                img = io.imread(filename)
                img = scipy.misc.imresize(img, imgSize)
                img = (img - MEAN_VALUES) / (255 / 2)
                img_batch.append(img)
            batch = np.concatenate(img_batch).reshape([-1, imgSize[0], imgSize[1], 3])
            return batch

        if self.cursor + self.batch_size > self.size:
            self.cursor = 0
            self.epoch += 1
        clear_batch = load_batch(
            [self.clear_path + '%d.jpg' % self.file_num[i] for i in range(self.cursor, self.cursor + self.batch_size)], img_size)
        blur_batch = load_batch(
            [self.blur_path + '%d.jpg' % self.file_num[i] for i in range(self.cursor, self.cursor + self.batch_size)], img_size)
        self.cursor += self.batch_size
        return clear_batch, blur_batch
